import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from numba import jit, njit, prange
from typing import List, Tuple, Dict, Optional
from collections import deque
import time
import pickle
import os


# 深度学习威胁评估模型
class ThreatAssessmentModel(nn.Module):
    """深度学习威胁评估模型"""

    def __init__(self, input_dim=9, hidden_dims=[64, 32, 16], output_dim=1):
        super(ThreatAssessmentModel, self).__init__()

        layers = []
        prev_dim = input_dim

        # 构建隐藏层
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(0.1))
            prev_dim = hidden_dim

        # 输出层
        layers.append(nn.Linear(prev_dim, output_dim))
        layers.append(nn.Sigmoid())  # 输出在0-1之间

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)


# 轻量级深度学习模型 (用于实时推理)
class FastThreatModel(nn.Module):
    """轻量级快速威胁评估模型"""

    def __init__(self, input_dim=9):
        super(FastThreatModel, self).__init__()

        self.network = nn.Sequential(
            nn.Linear(input_dim, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, 8),
            nn.ReLU(),
            nn.Linear(8, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.network(x)


# 在类外部定义Numba加速的函数
@njit(fastmath=True, cache=True)
def geodetic_to_enu_fast(lon: float, lat: float, alt: float,
                         ref_lon: float, ref_lat: float, ref_alt: float,
                         earth_radius: float, deg_to_rad: float) -> np.ndarray:
    """快速经纬高到ENU坐标转换"""
    ref_lat_rad = ref_lat * deg_to_rad

    e = (lon - ref_lon) * math.cos(ref_lat_rad) * earth_radius * math.pi / 180.0
    n = (lat - ref_lat) * earth_radius * math.pi / 180.0
    u = alt - ref_alt

    return np.array([e, n, u], dtype=np.float32)


@njit(fastmath=True, cache=True)
def estimate_velocity(positions: np.ndarray, times: np.ndarray) -> np.ndarray:
    """估计速度矢量"""
    if len(positions) < 2:
        return np.zeros(3, dtype=np.float32)

    pos1 = positions[-2]
    pos2 = positions[-1]
    time1 = times[-2]
    time2 = times[-1]

    if time2 <= time1:
        return np.zeros(3, dtype=np.float32)

    dt = time2 - time1
    velocity = (pos2 - pos1) / dt

    return velocity


@njit(fastmath=True, cache=True)
def calculate_trajectory_features(positions: np.ndarray, times: np.ndarray) -> np.ndarray:
    """计算轨迹特征 (用于深度学习模型输入)"""
    if len(positions) < 2:
        return np.zeros(9, dtype=np.float32)

    # 基本位置特征
    current_pos = positions[-1]
    
    # 手动计算均值 (避免使用axis参数)
    mean_pos = np.zeros(3, dtype=np.float32)
    for i in range(len(positions)):
        mean_pos[0] += positions[i, 0]
        mean_pos[1] += positions[i, 1]
        mean_pos[2] += positions[i, 2]
    mean_pos /= len(positions)
    
    # 手动计算标准差
    std_pos = np.zeros(3, dtype=np.float32)
    for i in range(len(positions)):
        std_pos[0] += (positions[i, 0] - mean_pos[0]) ** 2
        std_pos[1] += (positions[i, 1] - mean_pos[1]) ** 2
        std_pos[2] += (positions[i, 2] - mean_pos[2]) ** 2
    std_pos = np.sqrt(std_pos / len(positions))

    # 速度特征
    if len(positions) >= 2:
        velocities = []
        for i in range(1, len(positions)):
            dt = times[i] - times[i - 1]
            if dt > 1e-10:
                vel = (positions[i] - positions[i - 1]) / dt
                velocities.append(vel)

        if velocities:
            current_vel = velocities[-1]
            
            # 手动计算速度均值
            mean_vel = np.zeros(3, dtype=np.float32)
            for i in range(len(velocities)):
                mean_vel[0] += velocities[i][0]
                mean_vel[1] += velocities[i][1]
                mean_vel[2] += velocities[i][2]
            mean_vel /= len(velocities)
            
            # 手动计算速度标准差
            std_vel = np.zeros(3, dtype=np.float32)
            for i in range(len(velocities)):
                std_vel[0] += (velocities[i][0] - mean_vel[0]) ** 2
                std_vel[1] += (velocities[i][1] - mean_vel[1]) ** 2
                std_vel[2] += (velocities[i][2] - mean_vel[2]) ** 2
            std_vel = np.sqrt(std_vel / len(velocities))
        else:
            mean_vel = np.zeros(3, dtype=np.float32)
            std_vel = np.zeros(3, dtype=np.float32)
            current_vel = np.zeros(3, dtype=np.float32)
    else:
        mean_vel = np.zeros(3, dtype=np.float32)
        std_vel = np.zeros(3, dtype=np.float32)
        current_vel = np.zeros(3, dtype=np.float32)

    # 轨迹曲率特征
    curvature = 0.0
    if len(positions) >= 3:
        for i in range(1, len(positions) - 1):
            v1 = positions[i] - positions[i - 1]
            v2 = positions[i + 1] - positions[i]

            if np.linalg.norm(v1) > 1e-10 and np.linalg.norm(v2) > 1e-10:
                dot_product = v1[0] * v2[0] + v1[1] * v2[1] + v1[2] * v2[2]
                cos_angle = dot_product / (np.linalg.norm(v1) * np.linalg.norm(v2))
                cos_angle = max(min(cos_angle, 1.0), -1.0)
                angle = math.acos(cos_angle)
                curvature += angle

    # 组合所有特征
    features = np.zeros(9, dtype=np.float32)
    features[0:3] = current_pos  # 当前位置
    features[3:6] = current_vel  # 当前速度
    features[6] = curvature  # 轨迹曲率
    features[7] = np.linalg.norm(std_pos)  # 位置稳定性
    features[8] = np.linalg.norm(std_vel)  # 速度稳定性

    return features


class DeepLearningThreatEvaluator:
    """基于深度学习的威胁评估类"""

    # 类常量
    EARTH_RADIUS = 6371000.0
    DEG_TO_RAD = math.pi / 180.0

    def __init__(self, history_length: int = 10, max_stale_time: float = 5.0,
                 use_deep_learning: bool = True, model_path: str = None):
        """
        初始化威胁评估器

        参数:
            history_length: 历史轨迹长度
            max_stale_time: 最大过期时间(秒)
            use_deep_learning: 是否使用深度学习模型
            model_path: 预训练模型路径
        """
        # 基础参数
        self.history_length = history_length
        self.max_stale_time = max_stale_time
        self.use_deep_learning = use_deep_learning

        # 武器射程参考 (米)
        self.weapon_ranges = np.array([100000.0, 20000.0, 5000.0, 30000.0], dtype=np.float32)

        # 实体类型威胁系数
        self.entity_threats = np.array([1.0, 0.8, 0.4, 0.3, 0.9, 0.6], dtype=np.float32)

        # 历史轨迹存储
        self.trajectories = {}
        self.last_update_time = time.time()

        # 深度学习模型
        self.model = None
        self.feature_scaler = None

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if self.use_deep_learning:
            self._initialize_model(model_path)

        # 混合评估权重 (深度学习 + 规则基础)
        self.hybrid_weights = {
            'deep_learning': 0.7,
            'rule_based': 0.3
        }

        # 缓存系统
        self._threat_cache = {}
        self._cache_max_size = 1000
        self._feature_cache = {}

        # 性能监控
        self.inference_times = deque(maxlen=100)
        self.avg_inference_time = 0

        # 设备选择 (GPU优先)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {self.device}")

    def _initialize_model(self, model_path: str = None):
        """初始化深度学习模型"""
        if model_path and os.path.exists(model_path):
            # 加载预训练模型
            try:
                checkpoint = torch.load(model_path, map_location=self.device)
                self.model = FastThreatModel()
                self.model.load_state_dict(checkpoint['model_state_dict'])
                self.model.to(self.device)
                self.model.eval()  # 设置为评估模式

                # 加载特征缩放器
                if 'feature_scaler' in checkpoint:
                    self.feature_scaler = checkpoint['feature_scaler']

                print(f"成功加载预训练模型: {model_path}")
            except Exception as e:
                print(f"加载模型失败: {e}, 使用默认模型")
                self._create_default_model()
        else:
            # 创建默认模型
            self._create_default_model()

    def _create_default_model(self):
        """创建默认模型"""
        self.model = FastThreatModel()
        self.model.to(self.device)
        self.model.eval()

        # 创建默认特征缩放器
        self.feature_scaler = {
            'mean': np.zeros(9),
            'std': np.ones(9)
        }

        print("使用默认模型 (需要训练)")

    def update_trajectories(self, enemy_data: List[Tuple], timestamp: float = None):
        """更新敌机轨迹历史"""
        if timestamp is None:
            timestamp = time.time()

        # 清理过期敌机
        self._cleanup_stale_trajectories(timestamp)

        # 更新每个敌机的轨迹
        for plat_id, f_lon, f_lat, f_height in enemy_data:
            if plat_id not in self.trajectories:
                self.trajectories[plat_id] = {
                    'positions': deque(maxlen=self.history_length),
                    'times': deque(maxlen=self.history_length),
                    'enu_positions': deque(maxlen=self.history_length),
                    'last_updated': timestamp
                }

            # 存储原始位置和时间
            self.trajectories[plat_id]['positions'].append((f_lon, f_lat, f_height))
            self.trajectories[plat_id]['times'].append(timestamp)
            self.trajectories[plat_id]['last_updated'] = timestamp

            # 清除缓存
            if plat_id in self._threat_cache:
                del self._threat_cache[plat_id]
            if plat_id in self._feature_cache:
                del self._feature_cache[plat_id]

    def _cleanup_stale_trajectories(self, current_time: float):
        """清理过期的敌机轨迹"""
        stale_ids = []

        for plat_id, trajectory in self.trajectories.items():
            if current_time - trajectory['last_updated'] > self.max_stale_time:
                stale_ids.append(plat_id)

        for plat_id in stale_ids:
            del self.trajectories[plat_id]
            if plat_id in self._threat_cache:
                del self._threat_cache[plat_id]
            if plat_id in self._feature_cache:
                del self._feature_cache[plat_id]

    def _extract_features(self, plat_id: str, ally_data: Tuple) -> np.ndarray:
        """提取敌机特征 (用于深度学习模型)"""
        if plat_id in self._feature_cache:
            return self._feature_cache[plat_id]

        trajectory = self.trajectories[plat_id]
        ally_lon, ally_lat, ally_alt = ally_data

        # 确保有ENU坐标
        if len(trajectory['enu_positions']) != len(trajectory['positions']):
            trajectory['enu_positions'].clear()
            for lon, lat, alt in trajectory['positions']:
                enu_pos = geodetic_to_enu_fast(
                    lon, lat, alt, ally_lon, ally_lat, ally_alt,
                    self.EARTH_RADIUS, self.DEG_TO_RAD
                )
                trajectory['enu_positions'].append(enu_pos)

        # 转换为NumPy数组
        enu_positions = np.array(list(trajectory['enu_positions']), dtype=np.float32)
        times = np.array(list(trajectory['times']), dtype=np.float32)

        # 计算特征
        features = calculate_trajectory_features(enu_positions, times)

        # 缓存特征
        self._feature_cache[plat_id] = features

        return features

    def _deep_learning_threat_assessment(self, features: np.ndarray) -> float:
        """使用深度学习模型评估威胁"""
        if self.model is None:
            return 0.5  # 默认值

        start_time = time.time()

        try:
            # 特征预处理
            if self.feature_scaler is not None:
                features = (features - self.feature_scaler['mean']) / self.feature_scaler['std']

            # 转换为Tensor
            features_tensor = torch.FloatTensor(features).unsqueeze(0).to(self.device)

            # 模型推理
            with torch.no_grad():
                threat_tensor = self.model(features_tensor)
                threat = threat_tensor.cpu().numpy()[0, 0]

            # 记录推理时间
            inference_time = time.time() - start_time
            self.inference_times.append(inference_time)
            self.avg_inference_time = np.mean(self.inference_times) if self.inference_times else 0

            return float(threat)

        except Exception as e:
            print(f"深度学习推理失败: {e}")
            return 0.5

    def _rule_based_threat_assessment(self, plat_id: str, ally_data: Tuple,
                                      entity_type: int, weapon_type: int) -> float:
        """基于规则的威胁评估 (备用方法)"""
        trajectory = self.trajectories[plat_id]
        ally_lon, ally_lat, ally_alt = ally_data

        # 获取最新位置
        latest_lon, latest_lat, latest_alt = trajectory['positions'][-1]

        # 转换为ENU坐标
        enemy_pos_enu = geodetic_to_enu_fast(
            latest_lon, latest_lat, latest_alt,
            ally_lon, ally_lat, ally_alt,
            self.EARTH_RADIUS, self.DEG_TO_RAD
        )

        # 计算距离
        distance = np.linalg.norm(enemy_pos_enu)

        # 简单规则: 距离越近威胁越大
        weapon_range = self.weapon_ranges[weapon_type - 1]
        if distance > weapon_range * 1.2:
            return 0.0

        # 距离威胁 (指数衰减)
        distance_threat = math.exp(-0.0001 * distance)

        # 实体类型威胁
        entity_idx = min(entity_type - 1, len(self.entity_threats) - 1)
        entity_threat = self.entity_threats[entity_idx]

        return distance_threat * entity_threat

    def evaluate_threats(self, ally_data: Tuple, entity_types: Dict[str, int] = None,
                         weapon_types: Dict[str, int] = None) -> Dict[str, float]:
        """评估所有敌机的威胁 (混合方法)"""
        if entity_types is None:
            entity_types = {}
        if weapon_types is None:
            weapon_types = {}

        threats = {}

        for plat_id, trajectory in self.trajectories.items():
            if len(trajectory['positions']) < 2:
                continue

            # 检查缓存
            cache_key = f"{plat_id}_{ally_data}"
            if cache_key in self._threat_cache:
                threats[plat_id] = self._threat_cache[cache_key]
                continue

            # 获取实体类型和武器类型
            entity_type = entity_types.get(plat_id, 5)  # 默认UNKNOWN
            weapon_type = weapon_types.get(plat_id, 3)  # 默认CANNON

            # 混合威胁评估
            if self.use_deep_learning and self.model is not None:
                # 提取特征
                features = self._extract_features(plat_id, ally_data)

                # 深度学习评估
                dl_threat = self._deep_learning_threat_assessment(features)

                # 规则基础评估
                rb_threat = self._rule_based_threat_assessment(plat_id, ally_data, entity_type, weapon_type)

                # 混合评估
                threat = (self.hybrid_weights['deep_learning'] * dl_threat +
                          self.hybrid_weights['rule_based'] * rb_threat)
            else:
                # 仅使用规则基础评估
                threat = self._rule_based_threat_assessment(plat_id, ally_data, entity_type, weapon_type)

            # 应用武器射程限制
            weapon_range = self.weapon_ranges[weapon_type - 1]
            ally_lon, ally_lat, ally_alt = ally_data
            latest_lon, latest_lat, latest_alt = trajectory['positions'][-1]

            enemy_pos_enu = geodetic_to_enu_fast(
                latest_lon, latest_lat, latest_alt,
                ally_lon, ally_lat, ally_alt,
                self.EARTH_RADIUS, self.DEG_TO_RAD
            )

            distance = np.linalg.norm(enemy_pos_enu)
            if distance > weapon_range * 1.2:
                threat = 0.0

            threats[plat_id] = threat
            self._threat_cache[cache_key] = threat

            # 管理缓存大小
            if len(self._threat_cache) > self._cache_max_size:
                oldest_key = next(iter(self._threat_cache))
                del self._threat_cache[oldest_key]

        return threats

    def online_learning_update(self, new_data: List[Tuple], labels: List[float],
                               learning_rate: float = 0.001, batch_size: int = 32):
        """在线学习更新模型参数"""
        if not self.use_deep_learning or self.model is None:
            return

        try:
            # 转换为训练模式
            self.model.train()

            # 准备训练数据
            features_list = []
            for plat_id, ally_data in new_data:
                if plat_id in self.trajectories:
                    features = self._extract_features(plat_id, ally_data)
                    features_list.append(features)

            if len(features_list) < batch_size:
                return  # 数据不足

            # 转换为Tensor
            features_tensor = torch.FloatTensor(features_list[:batch_size]).to(self.device)
            labels_tensor = torch.FloatTensor(labels[:batch_size]).unsqueeze(1).to(self.device)

            # 优化器
            optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
            criterion = nn.MSELoss()

            # 训练步骤
            optimizer.zero_grad()
            outputs = self.model(features_tensor)
            loss = criterion(outputs, labels_tensor)
            loss.backward()
            optimizer.step()

            # 返回评估模式
            self.model.eval()

            print(f"在线学习更新完成, 损失: {loss.item():.4f}")

        except Exception as e:
            print(f"在线学习失败: {e}")
            self.model.eval()  # 确保返回评估模式

    def save_model(self, file_path: str):
        """保存模型和特征缩放器"""
        if self.model is None:
            return

        try:
            checkpoint = {
                'model_state_dict': self.model.state_dict(),
                'feature_scaler': self.feature_scaler,
                'hybrid_weights': self.hybrid_weights
            }

            torch.save(checkpoint, file_path)
            print(f"模型已保存到: {file_path}")

        except Exception as e:
            print(f"保存模型失败: {e}")

    def get_performance_stats(self) -> Dict:
        """获取性能统计信息"""
        return {
            'avg_inference_time_ms': self.avg_inference_time * 1000,
            'recent_inference_times': list(self.inference_times),
            'cache_size': len(self._threat_cache),
            'trajectory_count': len(self.trajectories),
            'device': str(self.device)
        }


# 使用示例
if __name__ == "__main__":
    # 创建评估器
    evaluator = DeepLearningThreatEvaluator(
        history_length=10,
        max_stale_time=3.0,
        use_deep_learning=True,
        model_path=None  # 无预训练模型，使用默认模型
    )

    # 模拟雷达数据
    ally_data = (116.4, 39.9, 10000)

    # 模拟多帧敌机数据
    enemy_frames = [
        [
            ("enemy_001", 116.42, 39.92, 10000),
            ("enemy_002", 116.38, 39.95, 12000),
        ],
        [
            ("enemy_001", 116.421, 39.921, 10000),
            ("enemy_002", 116.379, 39.951, 11900),
        ],
        [
            ("enemy_001", 116.422, 39.922, 10000),
            ("enemy_002", 116.378, 39.952, 11800),
        ]
    ]

    # 实体类型和武器类型映射
    entity_types = {"enemy_001": 1, "enemy_002": 2}
    weapon_types = {"enemy_001": 1, "enemy_002": 2}

    # 模拟处理多帧数据
    for i, frame in enumerate(enemy_frames):
        timestamp = i * 0.1
        evaluator.update_trajectories(frame, timestamp)

        # 评估威胁
        threats = evaluator.evaluate_threats(ally_data, entity_types, weapon_types)

        print(f"帧 {i + 1} 威胁评估:")
        for plat_id, threat in threats.items():
            print(f"  {plat_id}: {threat:.3f}")

        # 获取性能统计
        if i % 10 == 0:  # 每10帧输出一次性能统计
            stats = evaluator.get_performance_stats()
            print(f"性能统计: 平均推理时间 {stats['avg_inference_time_ms']:.2f}ms")

        print()

    # 性能测试
    start_time = time.time()

    # 模拟100次评估
    for _ in range(100):
        evaluator.evaluate_threats(ally_data, entity_types, weapon_types)

    end_time = time.time()
    print(f"100次评估耗时: {(end_time - start_time) * 1000:.2f} 毫秒")

    # 保存模型
    evaluator.save_model("threat_model.pth")